# models/database.py
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
import logging
import time
from functools import wraps

Base = declarative_base()

def retry_on_db_disconnect(max_retries=3, delay=1):
    """
    装饰器：在数据库连接断开时自动重试
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            retries = 0
            while retries < max_retries:
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    # 检查是否是连接相关的错误
                    if "MySQL server has gone away" in str(e) or "Can't connect to MySQL server" in str(e):
                        retries += 1
                        if retries >= max_retries:
                            raise e
                        logging.warning(f"⚠️ 数据库连接断开，第 {retries} 次重试...")
                        time.sleep(delay * retries)  # 指数退避
                        # 重新创建数据库引擎
                        if hasattr(args[0], 'engine'):
                            args[0].reconnect()
                    else:
                        raise e
            return func(*args, **kwargs)
        return wrapper
    return decorator

class DatabaseManager:
    def __init__(self, db_config):
        self.db_config = db_config
        self._create_engine()
        self._Session = sessionmaker(bind=self.engine)
        logging.info("✅ 数据库连接初始化成功")
    
    def _create_engine(self):
        """创建数据库引擎"""
        self.engine = create_engine(
            f"mysql+pymysql://{self.db_config['user']}:{self.db_config['password']}"
            f"@{self.db_config['host']}:{self.db_config['port']}/{self.db_config['database']}?charset={self.db_config.get('charset', 'utf8mb4')}",
            pool_recycle=3600,  # 1小时回收连接
            pool_pre_ping=True,  # 检查连接有效性
            echo=False
        )
    
    def reconnect(self):
        """重新连接数据库"""
        try:
            self.engine.dispose()
            self._create_engine()
            self._Session = sessionmaker(bind=self.engine)
            logging.info("✅ 数据库重新连接成功")
        except Exception as e:
            logging.error(f"❌ 数据库重新连接失败: {str(e)}")
            raise

    def Session(self):
        return self._Session()

    @retry_on_db_disconnect()
    def save_to_db(self, items, table_name=None):
        """通用保存函数"""
        from sqlalchemy.dialects.mysql import insert

        if not items:
            return

        if isinstance(items[0], dict):
            table_class = next((cls for cls in Base.__subclasses__() if cls.__tablename__ == table_name), None)
            if not table_class:
                raise ValueError(f"❌ 未找到对应的表类: {table_name}")

            insert_stmt = insert(table_class).values(items)
            on_duplicate_key_stmt = insert_stmt.on_duplicate_key_update(**{
                key: value for key, value in insert_stmt.inserted.items()
                if key not in {'id', 'created_at'}
            })

            with self.Session() as session:
                session.execute(on_duplicate_key_stmt)
                session.commit()
        else:
            with self.Session() as session:
                session.bulk_save_objects(items)
                session.commit()

    def close_connection(self):
        self.engine.dispose()